from ultralytics import YOLO
import cv2
model = YOLO(r'train7\weights\best.pt')  # load an official model
"""这个传入参数是路径，用于单个数字识别的"""
def Handwritten_digit_recognition_Single(savePath):
    img = cv2.imread(savePath)
    # 确保图片正确读取
    if img is not None:
    # 颜色反转（取反）  
    # OpenCV 使用 BGR 格式，因此我们需要对每个通道分别取反  
        img_inverted = 255 - img
    # 确保数据类型为 uint8，因为取反后可能超出 0-255 范围  
        img_inverted = img_inverted.astype('uint8')
    # 保存反转后的图片  
        savePathInverted = r"Single_inverted.PNG"
        cv2.imwrite(savePathInverted, img_inverted)
    # 打印消息以确认操作完成  
        print(f"反转图片保存到 {savePathInverted}")
    else:
        print("未读取到图片")

    # Predict with the model
    img = cv2.imread(savePathInverted, cv2.IMREAD_GRAYSCALE)
    results = model(savePathInverted,save=True)  # predict on an image
    return results
"""
# 对结果进行处理
for result in results:
 
    probs = result.probs
    print(probs)
"""
"""这个的传入参数是图片 用于连续数字识别的"""
def Handwritten_digit_recognition_coiled(img):
    if img is not None:
    # 颜色反转（取反）
    # OpenCV 使用 BGR 格式，因此我们需要对每个通道分别取反
        img_inverted = 255 - img
    # 确保数据类型为 uint8，因为取反后可能超出 0-255 范围
        img_inverted = img_inverted.astype('uint8')
    # 保存反转后的图片
        savePathInverted = r"inverted.PNG"
        cv2.imwrite(savePathInverted, img_inverted)
    # 打印消息以确认操作完成
        print(f"反转图片保存到 {savePathInverted}")
    else:
        print("未读取到图片")

    # Predict with the model
    img = cv2.imread(savePathInverted, cv2.IMREAD_GRAYSCALE)
    results = model(savePathInverted,save=True)  # predict on an image
    return results
